#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Q3 Method A: Solar Spectrum Tracking with 5-Channel LED (NNLS + Temporal Smoothing)
# Author: ChatGPT
# Python: 3.10+
# Dependencies: numpy, pandas, matplotlib, openpyxl
#
# Usage:
#   python q3_methodA.py --excel "/path/to/附录.xlsx" --lam 10 --iters 200 --outdir "./outputs"
#
# Outputs:
#   - Q3_MethodA_schedule.csv : per-time channel weights (sum=1), brightness, NRMSE
#   - nrmse_over_time.png     : NRMSE curve before/after smoothing
#   - spectrum_<time>.png     : Spectrum comparison at selected times (early/noon/late)

import argparse
import os
import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def extract_wavelength_series(col):
    """Extract numeric wavelength from strings like '380(mW/m2/nm)' -> 380.0"""
    def to_num(x):
        if isinstance(x, (int, float, np.integer, np.floating)):
            return float(x)
        s = str(x)
        m = re.match(r"(\d+\.?\d*)", s)
        if not m:
            m = re.search(r"(\d+\.?\d*)", s)
        return float(m.group(1)) if m else np.nan
    return col.map(to_num).astype(float)


def trapz_area(y, x):
    return np.trapz(y, x)


def project_to_simplex(v):
    """Project v onto the probability simplex {w >= 0, sum w = 1}"""
    v = np.asarray(v, dtype=float)
    n = v.size
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u) - 1
    ind = np.arange(1, n+1)
    cond = u - cssv/ind > 0
    if not np.any(cond):
        theta = cssv[-1] / n
    else:
        rho = np.where(cond)[0][-1]
        theta = cssv[rho] / (rho + 1.0)
    w = np.maximum(v - theta, 0.0)
    s = w.sum()
    if s == 0:
        w = np.ones_like(w) / n
    else:
        w /= s
    return w


def nnls_on_simplex(A, y, max_iter=200, step=None, seed=0):
    """Solve: min ||A w - y||^2 s.t. w in simplex (Projected Gradient)."""
    rng = np.random.default_rng(seed)
    m, n = A.shape
    w = np.ones(n) / n

    # Step size via power iteration (||A||_2^2)
    if step is None:
        v = rng.standard_normal(n)
        v /= (np.linalg.norm(v) + 1e-12)
        for _ in range(20):
            v = A.T @ (A @ v)
            nv = np.linalg.norm(v)
            if nv < 1e-12:
                break
            v /= nv
        L = np.linalg.norm(A @ v)**2
        if not np.isfinite(L) or L <= 0:
            L = 1.0
        step = 1.0/(2.0*L)

    for _ in range(max_iter):
        r = A @ w - y
        grad = 2.0 * (A.T @ r)
        w = project_to_simplex(w - step * grad)
    return w


def smooth_optimize(A, S, wl, W_init, B_init, lam=10.0, n_iter=200, lamB=0.0):
    """Alternating optimization with temporal smoothing on W and closed-form B."""
    T = S.shape[1]
    n_ch = A.shape[1]
    W = W_init.copy()
    B = B_init.copy()

    AtA = A.T @ A
    ATS = A.T @ S

    # Step size for W (Lipschitz upper bound estimate)
    v = np.random.randn(n_ch)
    v /= (np.linalg.norm(v) + 1e-12)
    for _ in range(20):
        v = AtA @ v
        nv = np.linalg.norm(v)
        if nv < 1e-12:
            break
        v /= nv
    eig_approx = np.linalg.norm(AtA @ v)
    L_A = eig_approx if np.isfinite(eig_approx) and eig_approx > 0 else 1.0
    step_w = 1.0/(2*(np.max(B)**2)*L_A + 4*lam + 1e-6)

    for _ in range(n_iter):
        # Update W with one PGD step per time
        W_new = W.copy()
        for t in range(T):
            g = 2*(B[t]**2)*(AtA @ W[t]) - 2*B[t]*ATS[:, t]
            if t > 0:
                g += 2*lam*(W[t] - W[t-1])
            if t < T-1:
                g += 2*lam*(W[t] - W[t+1])
            w = W[t] - step_w*g
            W_new[t] = project_to_simplex(w)
        W = W_new
        # Update B (closed form, >=0)
        for t in range(T):
            Aw = A @ W[t]
            denom = np.dot(Aw, Aw)
            num = np.dot(Aw, S[:, t])
            B[t] = max(0.0, num/denom) if denom > 0 else 0.0
            # Optional: temporal smoothing on brightness B (quadratic penalty on slope)
        if lamB > 0:
            B_sm = B.copy()
            alpha = 0.1 * lamB  # conservative step to avoid over-smoothing
            for t in range(1, T-1):
                # gradient of (B[t]-B[t-1])^2 + (B[t]-B[t+1])^2
                gradBt = 2*(B[t]-B[t-1]) + 2*(B[t]-B[t+1])
                B_sm[t] = max(0.0, B[t] - alpha*gradBt)
            B = B_sm
    return W, B


def reconstruct_spd(A, W, B):
    T = W.shape[0]
    R = np.zeros((A.shape[0], T))
    for t in range(T):
        R[:, t] = B[t] * (A @ W[t])
    return R


def nrmse(S, R):
    num = np.sqrt(((S - R)**2).sum(axis=0))
    den = np.sqrt((S**2).sum(axis=0)) + 1e-12
    return num / den


def pick_times(time_cols, targets=("08:30:00", "12:30:00", "17:30:00")):
    chosen = []
    for tar in targets:
        if tar in time_cols:
            chosen.append(tar)
        else:
            # closest by minutes difference
            def to_min(s):
                h, m, s2 = map(int, s.split(":"))
                return 60*h + m + s2/60.0
            tm = to_min(tar)
            mins = np.array([to_min(x) for x in time_cols])
            idx = int(np.argmin(np.abs(mins - tm)))
            chosen.append(time_cols[idx])
    return chosen


def main():
    parser = argparse.ArgumentParser(
        description="Q3 Method A: Solar spectrum tracking via 5-channel LED")
    parser.add_argument("--excel", type=str, required=True,
                        help="Path to 附录.xlsx")
    parser.add_argument("--lam", type=float, default=10.0,
                        help="Temporal smoothing weight (lambda)")
    parser.add_argument("--iters", type=int, default=200,
                        help="Smoothing iterations")
    parser.add_argument("--lamB", type=float, default=0.0,
                        help="Temporal smoothing weight for brightness B (lambda_B)")
    parser.add_argument("--outdir", type=str,
                        default="./outputs", help="Directory to save outputs")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    # Load sheets
    sun_df = pd.read_excel(args.excel, sheet_name="Problem 3 SUN_SPD")
    led_df = pd.read_excel(args.excel, sheet_name="Problem 2_LED_SPD")

    # Common wavelengths
    wl_sun = extract_wavelength_series(sun_df.iloc[:, 0])
    wl_led = extract_wavelength_series(led_df.iloc[:, 0])
    common_wl = np.intersect1d(wl_sun.values, wl_led.values)

    sun = sun_df[sun_df.iloc[:, 0].apply(lambda x: float(re.match(
        r"(\d+\.?\d*)", str(x)).group(1)) if re.match(r"(\d+\.?\d*)", str(x)) else np.nan).isin(common_wl)].copy()
    led = led_df[led_df.iloc[:, 0].apply(lambda x: float(re.match(
        r"(\d+\.?\d*)", str(x)).group(1)) if re.match(r"(\d+\.?\d*)", str(x)) else np.nan).isin(common_wl)].copy()

    # Sort and set wavelength
    sun["wl"] = extract_wavelength_series(sun.iloc[:, 0])
    led["wl"] = extract_wavelength_series(led.iloc[:, 0])
    sun = sun.sort_values("wl").reset_index(drop=True)
    led = led.sort_values("wl").reset_index(drop=True)
    wl = sun["wl"].values

    def _norm_time_label(c):
        # 统一列名成 "HH:MM:SS" 字符串
        # 1) pandas 时间戳
        if isinstance(c, pd.Timestamp):
            return c.strftime("%H:%M:%S")
        # 2) 数字（Excel 里时间可能是“一天的小数”）
        try:
            v = float(c)
            if 0 <= v < 2:  # 允许 0~1（一天内），个别表会到 ~1.x
                sec = int(round(v * 24 * 3600))
                return f"{sec//3600:02d}:{(sec%3600)//60:02d}:{sec%60:02d}"
        except Exception:
            pass
        # 3) 普通字符串，补齐到 HH:MM:SS
        s = str(c).strip()
        if re.match(r"^\d{1,2}:\d{2}(:\d{2})?$", s):
            if s.count(":") == 1:
                s = s + ":00"
            return s
        return s  # 兜底：保持原样

    # 先找出“候选时间列”（除去第一列和 wl）
    _raw_time_cols = [
        c for c in sun.columns if c not in [sun.columns[0], "wl"]]
    # 把列名规范化
    _new_labels = {_c: _norm_time_label(_c) for _c in _raw_time_cols}
    sun.rename(columns=_new_labels, inplace=True)
    # 真正的时间列名

    def _to_min(s):
        h, m, s2 = map(int, str(s).split(":"))
        return h*60 + m + s2/60.0

    time_cols = sorted(list({_new_labels[c]
                       for c in _raw_time_cols}), key=_to_min)

    # Channel columns
    alias_map = {
        "Blue": ["Blue", "B"],
        "Green": ["Green", "G"],
        "Red": ["Red", "Deep Red", "DR"],
        "Warm White": ["Warm White", "WW", "WarmWhite"],
        "Cold White": ["Cold White", "Cool White", "CW", "CoolWhite", "ColdWhite"],
    }
    found = {}
    for std, aliases in alias_map.items():
        for a in aliases:
            if a in led.columns:
                found[std] = a
                break
        if std not in found:
            raise ValueError(f"LED SPD缺少通道：{std}（可用别名：{aliases}）")

    # 固定标准顺序，后续 W 的列顺序与之对齐
    channel_cols = ["Blue", "Green", "Red", "Warm White", "Cold White"]
    A = led[[found[c] for c in channel_cols]].values.astype(float)  # (n_wl, 5)

    # Normalize each channel by area
    areas = np.array([trapz_area(A[:, i], wl) for i in range(A.shape[1])])
    areas[areas == 0] = 1.0
    A_norm = A / areas

    # Solar SPDs
    S = sun[time_cols].values.astype(float)  # (n_wl, T)
    T = S.shape[1]

    # Initial per-time NNLS-on-simplex for shape
    sun_area = np.array([trapz_area(S[:, t], wl) for t in range(T)])
    sun_area[sun_area == 0] = 1.0
    S_hat = S / sun_area

    W0 = np.zeros((T, A_norm.shape[1]))
    B0 = np.zeros(T)
    for t in range(T):
        y = S_hat[:, t]
        W0[t] = nnls_on_simplex(A_norm, y, max_iter=200)
        Aw = A_norm @ W0[t]
        denom = np.dot(Aw, Aw)
        num = np.dot(Aw, S[:, t])
        B0[t] = max(0.0, num/denom) if denom > 0 else 0.0

    # Temporal smoothing
    W, B = smooth_optimize(A_norm, S, wl, W0, B0,
                           lam=args.lam, n_iter=args.iters, lamB=args.lamB)

    # Reconstructions and errors
    R0 = reconstruct_spd(A_norm, W0, B0)
    R = reconstruct_spd(A_norm, W,  B)
    nrmse0 = nrmse(S, R0)
    nrmse1 = nrmse(S, R)

    # Save schedule
    schedule = pd.DataFrame({
        "Time": time_cols,
        "Blue": W[:, channel_cols.index("Blue")],
        "Green": W[:, channel_cols.index("Green")],
        "Red": W[:, channel_cols.index("Red")],
        "Warm White": W[:, channel_cols.index("Warm White")],
        "Cold White": W[:, channel_cols.index("Cold White")],
        "Brightness": B,
        "NRMSE": nrmse1
    })
    sch_path = os.path.join(args.outdir, "Q3_MethodA_schedule.csv")
    schedule.to_csv(sch_path, index=False)
    print(f"[Saved] {sch_path}")

    # Plot NRMSE
    plt.figure()
    plt.plot(range(T), nrmse0, label="Before smoothing")
    plt.plot(range(T), nrmse1, label="After smoothing")
    plt.xlabel("Time index")
    plt.ylabel("NRMSE")
    plt.title("Spectral Fit Error Over Time")
    plt.legend()
    plt.tight_layout()
    fig1 = os.path.join(args.outdir, "nrmse_over_time.png")
    plt.savefig(fig1, dpi=200)
    print(f"[Saved] {fig1}")

    # 早晨/正午/傍晚（若无精确时刻会就近匹配；傍晚=最后一个时间点）
    chosen_times = pick_times(
        time_cols, ("08:30:00", "12:30:00", time_cols[-1]))
    for tl in chosen_times:
        idx = time_cols.index(tl)
        plt.figure()
        plt.plot(wl, S[:, idx], label=f"Sun {tl}")
        plt.plot(wl, R[:, idx], label=f"LED mix {tl}")
        plt.xlabel("Wavelength (nm)")
        plt.ylabel("Spectral Power (as given)")
        plt.title(f"Spectrum Comparison at {tl}")
        plt.legend()
        plt.tight_layout()
        outp = os.path.join(args.outdir, f"spectrum_{tl.replace(':','-')}.png")
        plt.savefig(outp, dpi=200)
        print(f"[Saved] {outp}")

    # Save preview of first 10 rows
    head_path = os.path.join(args.outdir, "Q3_MethodA_schedule_head.csv")
    schedule.head(10).to_csv(head_path, index=False)
    print(f"[Saved] {head_path}")


if __name__ == "__main__":
    main()
